#coding=utf-8
import os
import time
import numpy as np
import h5py
import argparse
import dm_env
import cv2
import queue

import rospy
from sensor_msgs.msg import JointState
from sensor_msgs.msg import Image
from nav_msgs.msg import Odometry
from cv_bridge import CvBridge

import threading
import collections

import sys
sys.path.append("./")



# 保存数据函数
def save_data(args, timesteps, actions, dataset_path):
    # 数据字典
    data_dict = {
        # obs包含qpos，qvel， effort ,acition, image, base_action
        '/observations/qpos': [],
        '/observations/qvel': [],
        '/observations/effort': [],
        '/action': [],
        # '/base_action': [],
    }

    while actions:
        # 循环弹出一个队列
        action = actions.pop(0)   # 动作  当前动作
        ts = timesteps.pop(0)     # 奖励  前一帧

        # 往字典里面添值
        data_dict['/action'].append(action)

    t0 = time.time()
    with h5py.File(dataset_path + '.hdf5', 'w', rdcc_nbytes=1024**2*2) as root:
        # 文本的属性：
        # 1 是否仿真
        # 2 图像是否压缩
        #
        root.attrs['sim'] = False
        root.attrs['compress'] = False
        if args.is_compress:
            root.attrs['compress'] = True

        # 创建一个新的组observations，观测状态组
        # 图像组
        obs = root.create_group('observations')
        
        image = obs.create_group('images')

        depth = obs.create_group('depths')

        _ = root.create_dataset('action', (args.max_timesteps, 6))

        for name, array in data_dict.items():   # 名字+值
            root[name][...] = array

    print(f'Saving: {time.time() - t0:.1f} secs', dataset_path)


class RosOperator:
    def __init__(self, args):
        self.init(args)

    def init(self, args):
        self.args = args
        self.masterl_queue = queue.Queue()

        dataset_dir = os.path.join(args.dataset_dir, args.task_name)
        if not os.path.exists(dataset_dir):
            os.makedirs(dataset_dir)
        self.dataset_path = os.path.join(dataset_dir, "episode_" + str(args.episode_idx))

        self.register_sub()   # 定义订阅消息话题

    def register_sub(self):
        rospy.init_node('record_episodes', anonymous=True)

        rospy.Subscriber(self.args.master_arm_left_topic,  JointState, lambda msg: self.masterl_queue.put(msg), queue_size=200, tcp_nodelay=True)

    def collet_data(self):
        rate = rospy.Rate(self.args.frame_rate)  # 设置数据收集频率
        input("\033[32m\nPlease press any key to continue\n\033[0m")
        
        count = 0  # 统计 timesteps
        timesteps = []  # obs
        actions = []  # action
        
        while (count < self.args.max_timesteps and not rospy.is_shutdown()):
            time1 = time.time()  # 起始时间

            # 检查队列是否为空
            if self.masterl_queue.empty():
                print("\033[31mPlease check topic\n\033[0m")
                rate.sleep()
                continue

            # 获取最小时间戳
            min_timestamps = self.masterl_queue.queue[-1].header.stamp.to_sec() - 0.01

            # 获取主臂数据
            while (self.masterl_queue.queue[0].header.stamp.to_sec() < min_timestamps):
                self.masterl_queue.get()
            masterl = self.masterl_queue.get()

            count += 1

            # 2.1 从臂的信息
            obs = {
                'qpos': np.zeros(6),  # 假设 qpos 是一个长度为 6 的数组
                'qvel': np.zeros(6),  # 假设 qvel 是一个长度为 6 的数组
                'effort': np.zeros(6),  # 假设 effort 是一个长度为 6 的数组
            }


            # 2.2 时间步
            if count == 1:
                ts = dm_env.TimeStep(
                    step_type=dm_env.StepType.FIRST,
                    reward=None,
                    discount=None,
                    observation=obs)
            else:
                ts = dm_env.TimeStep(
                    step_type=dm_env.StepType.MID,
                    reward=None,
                    discount=None,
                    observation=obs)

            # 2.3 动作
            action = np.array(masterl.position)
            actions.append(action)
            timesteps.append(ts)

            if rospy.is_shutdown():
                exit(-1)
            
            print("frame %s cost_time: %s" % (count, round((time.time() - time1), 6)))
            rate.sleep()
            
        print("len(timesteps): ", len(timesteps))
        print("len(actions)  : ", len(actions))
        save_data(self.args, timesteps, actions, self.dataset_path)

def get_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset_dir', action='store', type=str, help='Dataset_dir.',
                        default="./data", required=False)
    
    parser.add_argument('--task_name', action='store', type=str, help='Task name.',
                        default="aloha_mobile_dummy", required=False)
    
    parser.add_argument('--episode_idx', action='store', type=int, help='Episode index.',
                        default=0, required=False)
    parser.add_argument('--max_timesteps', action='store', type=int, help='Max_timesteps.',
                        default=500, required=False)

   
    parser.add_argument('--master_arm_left_topic', action='store', type=str, help='master_arm_left_topic',
                        default='/master/joint_left', required=False)
    parser.add_argument('--frame_rate', action='store', type=int, help='frame_rate',
                        default=20, required=False)
    args = parser.parse_args()
    return args


def main():
    args = get_arguments()
    ros_operator = RosOperator(args)   # 初始化ros节点，订阅消息

    t1 = threading.Thread(target=ros_operator.collet_data)
    t1.start()
    t1.join()
    

if __name__ == '__main__':
    main()

# python scripts/collect_master.py
